package lzc.encrypt.encryption;

import com.github.davidmoten.rtree.geometry.Rectangle;
import com.google.common.hash.BloomFilter;
import com.google.common.hash.Funnels;
import it.unisa.dia.gas.jpbc.Element;
import it.unisa.dia.gas.jpbc.ElementPowPreProcessing;
import it.unisa.dia.gas.jpbc.PairingPreProcessing;
import lzc.encrypt.Assit;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;

/**
 * 参数类，IPPE使用到的参数
 *
 * @author 蓝色的三明治
 * @date 2024/09/20
 */
public class Parameter {
    public static double fpp = 0.001;// 布隆过滤器的容错率

    private PairingParameter pairingParameter;// 双线性映射相关参数

    private BigInteger max;// 向量计算的最大值

    private BloomFilter<String> bloomFilter;// 用于检测加密向量计算结果的布隆过滤器

    private EncryptionKey dataKey, searchKey;// 用于加密数据向量的密钥组 dataKey 和加密搜索向量的密钥组 searchKey

    private BigInteger randForRect;// 用于加密矩形区域的随机值

    /**
     * @param dimension 向量维度
     * @param max       内积的最大值
     */
    public Parameter(int dimension, BigInteger max) {
        // 初始化双线性映射与布隆过滤器
        this.pairingParameter = new PairingParameter();
        this.max = max;

        // 生成用于加密的密钥组
        this.searchKey = new EncryptionKey(pairingParameter.getG(), pairingParameter.getFromZp(), pairingParameter.getP(), dimension);
        this.dataKey = new EncryptionKey(pairingParameter.getG(), pairingParameter.getFromZp(), searchKey.getKeyMatrix());

        // 根据生成内容构建布隆过滤器
        // 将{H(e(g,g)^(abi))}_{i=1}^{max^2+max}插入到布隆过滤器中，其中H()是一种安全哈希函数，这里使用SHA-256
        BigInteger upper_limit_number = max.pow(2).add(max.sqrt().multiply(BigInteger.TWO));
//        BigInteger upper_limit_number = max.multiply(max.sqrt()).add(max.sqrt().sqrt().multiply(BigInteger.TWO));
        this.bloomFilter = BloomFilter.create(Funnels.stringFunnel(StandardCharsets.UTF_8), upper_limit_number.longValue(), fpp);
        Element upper_limit = pairingParameter.getFromZp(upper_limit_number).getImmutable();
        Element EGG = pairingParameter.pairing(pairingParameter.getG(), pairingParameter.getG()).getImmutable();// 计算 e(g,g)
        Element ab = dataKey.getN().mul(searchKey.getN()).getImmutable();// 计算 ab
        Element EGG_ab = EGG.powZn(ab).getImmutable();// 计算 e(g,g)^ab
        ElementPowPreProcessing EGG_ab_pre = EGG_ab.getElementPowPreProcessing();// 由于需要多次计算，为了提高效率，使用预处理元素
        for (Element i = pairingParameter.getFromZp(0); !i.isEqual(upper_limit); i.add(pairingParameter.getFromZp(1))) {
            Element value = EGG_ab_pre.powZn(i).getImmutable();// 计算 e(g,g)^(abi)
            String hash = Assit.getHash(value.toBytes());
            bloomFilter.put(hash);// 将hash值插入到布隆过滤器中
        }

        randForRect = pairingParameter.getFromZp().toBigInteger();
    }

    /**
     * 加密数据向量
     *
     * @param dataVector 明文数据向量
     * @return {@link DataVector} 加密完成的数据向量
     */
    public DataVector encryptDataVector(BigInteger[] dataVector) {
        // 获取加密完成的数据向量
        Element[] vector3 = encrypt(dataVector, true);

        // 数据向量要长期存储，并用于双线性映射计算，为了提高计算效率，将Element转化为PairingPreProcessing
        PairingPreProcessing[] vector4 = new PairingPreProcessing[vector3.length];
        for (int i = 0; i < vector3.length; i++) {
            vector4[i] = pairingParameter.pairing(vector3[i]);
        }
        return new DataVector(vector4);
    }

    /**
     * 加密矩形范围，生成两个对应该范围的数据向量
     *
     * @param rectangle 矩形范围
     * @return {@link DataVector[]} 两个对应该范围的数据向量
     */
    public DataVector[] encryptDataVector(Rectangle rectangle) {
        // 获取加密完成的数据向量
        Element[][] vector3 = encrypt(rectangle, true);

        // 数据向量要长期存储，并用于双线性映射计算，为了提高计算效率，将Element转化为PairingPreProcessing
        PairingPreProcessing[] vector4_1 = new PairingPreProcessing[vector3[0].length];
        PairingPreProcessing[] vector4_2 = new PairingPreProcessing[vector3[0].length];
        for (int i = 0; i < vector3[0].length; i++) {
            vector4_1[i] = pairingParameter.pairing(vector3[0][i]);
            vector4_2[i] = pairingParameter.pairing(vector3[1][i]);
        }
        return new DataVector[]{new DataVector(vector4_1), new DataVector(vector4_2)};
    }

    /**
     * 加密搜索向量
     *
     * @param searchVector 明文搜索向量
     * @return {@link SearchVector} 加密完成的搜索向量
     */
    public SearchVector encryptSearchVector(BigInteger[] searchVector) {
        return new SearchVector(encrypt(searchVector, false));
    }

    /**
     * 加密矩形范围，生成两个对应该范围的搜索向量
     *
     * @param rectangle 矩形范围
     * @return {@link SearchVector[]} 两个对应该范围的搜索向量
     */
    public SearchVector[] encryptSearchVector(Rectangle rectangle) {
        // 获取加密完成的搜索向量
        Element[][] vector3 = encrypt(rectangle, false);

        return new SearchVector[]{new SearchVector(vector3[0]), new SearchVector(vector3[1])};
    }

    /**
     * 加密矩形范围
     *
     * @param rectangle 矩形范围
     * @return {@link DataVector[]} 对应矩形范围的两个加密向量
     */
    private Element[][] encrypt(Rectangle rectangle, boolean ifData) {
        // 获取Rectangle对应的范围，并将其转化为整数范围
        BigInteger minX = BigInteger.valueOf((int) Math.floor(rectangle.x1()));
        BigInteger maxX = BigInteger.valueOf((int) Math.ceil(rectangle.x2()));

        // 根据范围转化为向量
        Element[] vectorRange1 = encrypt(minX, maxX, ifData);

        BigInteger minY = BigInteger.valueOf((int) Math.floor(rectangle.y1()));
        BigInteger maxY = BigInteger.valueOf((int) Math.ceil(rectangle.y2()));

        // 根据范围转化为向量
        Element[] vectorRange2 = encrypt(minY, maxY, ifData);

        return new Element[][]{vectorRange1, vectorRange2};
    }

    /**
     * 将范围转化为加密的整数向量
     *
     * @param min    范围下限
     * @param max    范围上限
     * @param ifData true表示明文向量是数据向量，false表示是搜索向量
     * @return {@link DataVector} 对应范围的加密向量
     */
    private Element[] encrypt(BigInteger min, BigInteger max, boolean ifData) {
        // 将范围转化为整数向量
        BigInteger[] vector;
        if (ifData) {
            vector = new BigInteger[]{randForRect.subtract(min.multiply(max)), BigInteger.ONE, min, max};
        } else {
            vector = new BigInteger[]{BigInteger.ONE, randForRect.negate().subtract(min.multiply(max)), min, max};
        }

        // 加密整数向量
        return encrypt(vector, ifData);
    }

    /**
     * 加密向量
     *
     * @param vector 明文向量
     * @param ifData true表示明文向量是数据向量，false表示是搜索向量
     * @return {@link Element[]} 加密完成的向量
     */
    private Element[] encrypt(BigInteger[] vector, boolean ifData) {
        // 随机获取y1，y2的值，满足 |y1| < y2 <= max^1/4
        BigInteger y1, y2;
        BigInteger upper_limit = max.sqrt();
        y2 = Assit.getRandom(BigInteger.valueOf(3), upper_limit);
        y1 = Assit.getRandom(BigInteger.ZERO, y2);

        BigInteger[] vector1 = new BigInteger[vector.length + 2];
        for (int i = 0; i < vector.length; i++) {
            vector1[i] = vector[i].multiply(y2);
        }

        if (ifData) {
            // 将进行过一次加密的向量进行二次加密
            vector1[vector1.length - 2] = y1;
            vector1[vector1.length - 1] = BigInteger.ONE;
            return dataKey.encrypt(vector1, pairingParameter);
        } else {
            vector1[vector1.length - 2] = BigInteger.ONE;
            vector1[vector1.length - 1] = y1;
            return searchKey.encrypt(vector1, pairingParameter);
        }
    }

    /**
     * 加密情况下，检查两个向量之间的内积是否不小于 0
     *
     * @param vector1 第一个用于计算的向量
     * @param vector2 第二个用于计算的向量
     * @return boolean true表示内积不小于 0，false表示内积小于 0
     */
    public boolean check(Element[] vector1, Element[] vector2) {
        if (vector1.length != vector2.length) {
            throw new RuntimeException("计算的两个向量长度不相等");
        }

        // 计算e(A,B)^{u" * v"}
        Element innerProduct = pairingParameter.pairing(vector1, vector2);

        // 获取计算结果的哈希值，若该值存在于布隆过滤器中，则内积结果不小于 0，反之则小于 0
        String hash = Assit.getHash(innerProduct.toBytes());
        return bloomFilter.mightContain(hash);
    }

    /**
     * 加密情况下，检查两个向量之间的内积是否不小于 0
     *
     * @param vector1 第一个用于计算的向量
     * @param vector2 第二个用于计算的向量
     * @return boolean true表示内积不小于 0，false表示内积小于 0
     */
    public boolean check(PairingPreProcessing[] vector1, Element[] vector2) {
        if (vector1.length != vector2.length) {
            throw new RuntimeException("计算的两个向量长度不相等");
        }

        // 计算e(A,B)^{u" * v"}
        Element innerProduct = pairingParameter.getG2().newElement(1);
        for (int i = 0; i < vector1.length; i++) {
            innerProduct.mul(vector1[i].pairing(vector2[i]));
        }

        // 获取计算结果的哈希值，若该值存在于布隆过滤器中，则内积结果不小于 0，反之则小于 0
        String hash = Assit.getHash(innerProduct.toBytes());
        return bloomFilter.mightContain(hash);
    }

    /**
     * 加密情况下，检查两个向量之间的内积是否不小于 0
     *
     * @param vector1 第一个用于计算的向量
     * @param vector2 第二个用于计算的向量
     * @return boolean true表示内积不小于 0，false表示内积小于 0
     */
    public boolean check(DataVector vector1, SearchVector vector2) {
        if (vector1.length != vector2.length) {
            throw new RuntimeException("计算的两个向量长度不相等");
        }

        // 计算e(A,B)^{u" * v"}
        Element innerProduct = pairingParameter.getG2().newElement(1);
        for (int i = 0; i < vector1.length; i++) {
            innerProduct.mul(vector1.getFromDataVector(i).pairing(vector2.getFromSearchVector(i)));
        }

        // 获取计算结果的哈希值，若该值存在于布隆过滤器中，则内积结果不小于 0，反之则小于 0
        String hash = Assit.getHash(innerProduct.toBytes());
        return bloomFilter.mightContain(hash);
    }

    /**
     * 加密情况下，检查两组向量之间的内积是否不小于 0
     *
     * @param vector1 第一个用于计算的向量
     * @param vector2 第二个用于计算的向量
     * @return boolean true表示内积不小于 0，false表示内积小于 0
     */
    public boolean check(DataVector[] vector1, SearchVector[] vector2) {
        if (vector1.length != vector2.length) {
            throw new IllegalArgumentException("两组向量的数量必须相等！");
        }

        // 逐个检查内积是否都不小于0，过检查到最后都不小于0，返回true，否则停止检查，返回false
        for (int i = 0; i < vector1.length; i++) {
            if (!check(vector1[i], vector2[i])) {
                return false;
            }
        }

        return true;
    }

    public EncryptionKey getDataKey() {
        return dataKey;
    }

    public EncryptionKey getSearchKey() {
        return searchKey;
    }

    public PairingParameter getPairingParameter() {
        return pairingParameter;
    }
}